package com.witmore.neutron.ai.cognitive.openai;

import com.alibaba.fastjson.JSONArray;
import com.witmore.neutron.ai.api.dto.ContentDTO;
import com.witmore.neutron.ai.api.dto.MessageDTO;
import com.witmore.neutron.ai.api.enums.ChatClientType;
import com.witmore.neutron.ai.api.service.ChatClientAdapter;
import com.witmore.neutron.ai.cognitive.ChatFluxUtilSSe;
import com.witmore.neutron.ai.cognitive.advisor.CustomSessionUsageAdvisor;
import com.witmore.neutron.ai.cognitive.service.AgentPluginService;
import com.witmore.neutron.ai.cognitive.service.UserSessionUsageService;
import com.witmore.neutron.ai.db.entity.ModelConfigDO;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.ChatModelStreamAdvisor;
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.content.Media;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.stereotype.Service;
import org.springframework.util.MimeType;
import org.springframework.util.MimeTypeUtils;
import org.springframework.util.ObjectUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.client.ResourceAccessException;
import reactor.core.publisher.Flux;

import java.net.URI;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Description: MultiModelService
 * Author: 豆子高
 * Date: 2025/5/30 10:09
 */
@Slf4j
@Service
public class OpenAiChatClientService implements ChatClientAdapter {

    /**
     * 在服务类中初始化ChatMemory
     */
    private final Map<String, ChatMemory> sessionChatMemoryMap = new ConcurrentHashMap<>();

    /**
     * 默认系统提示词
     */
    private static final String DEFAULT_SYSTEM_PROMPT = "你是一个中立、简洁、有用的AI助手。（不知道的，不要随意回答，就委婉的说你不知道））";

    /**
     * 默认用户消息
     */
    private static final String DEFAULT_USER_MESSAGE = "（用户未提供输入）";

    @Autowired
    private AiClientModelCache aiClientModelCache;

    @Autowired
    private ChatMemoryRepository chatMemoryRepository;

    @Autowired
    private AgentPluginService agentPluginService;

    @Autowired
    private UserSessionUsageService userSessionUsageService;

    @Override
    public String generateResponse(ModelConfigDO modelConfigDO, MessageDTO messageBO) {
        log.info("Generating response using OpenAI model: {}", modelConfigDO.getModelName());
        ChatMemory chatMemory = sessionChatMemoryMap.computeIfAbsent(messageBO.getSessionId(), sessionId -> MessageWindowChatMemory.builder().chatMemoryRepository(chatMemoryRepository).maxMessages(messageBO.getMaxMessages() == null ? 10 : messageBO.getMaxMessages()).build());

        OpenAiChatModel chatModelFromCache = this.getChatModelFromCache(modelConfigDO.getId().toString());
        Prompt prompt = this.buildChatModelRequest(chatMemory, messageBO);
//        log.info("Generated prompt for OpenAI model: {}", prompt);
        try {
            ChatClient chatClient = ChatClient.builder(chatModelFromCache)
                    .defaultAdvisors(
                            PromptChatMemoryAdvisor
                                    .builder(chatMemory)
                                    .conversationId(messageBO.getSessionId())
                                    .order(1)
                                    .build(),
                            new CustomSessionUsageAdvisor(userSessionUsageService)
                    )
                    .build();

            ToolCallback[] toolCallbacks = agentPluginService.getToolCallbacks(messageBO.getSessionId());
            Map<String, Object> context = getContextMap(modelConfigDO, messageBO, toolCallbacks);

            // 构建ChatClient 注册MCP
            ChatResponse chatResponse = Optional.ofNullable(agentPluginService.getToolCallbacks(messageBO.getSessionId()))
                    .filter(callbacks -> callbacks.length > 0)
                    .map(callbacks -> chatClient.prompt(prompt)
                            .advisors(advisorSpec -> advisorSpec.params(context))
                            .toolCallbacks(callbacks)
                            .call())
                    .orElseGet(() -> chatClient.prompt(prompt)
                            .advisors(advisorSpec -> advisorSpec.params(context))
                            .call()).chatResponse();
            Usage usage = Objects.requireNonNull(chatResponse).getMetadata().getUsage();
            log.info("Usage: {}", usage);
            AssistantMessage output = Objects.requireNonNull(chatResponse).getResult().getOutput();
            Map<String, Object> metadata = output.getMetadata();
            metadata.put("usage", usage);
            String content = output.getText();
            log.info("generateResponse Generated generateResponse: {}", output);
            chatMemory.add(messageBO.getSessionId(), output);
            return content;
        } catch (ResourceAccessException e) {
            // TODO 此处细化异常 日后可能会针对网络异常进行重试等处理
            chatMemory.add(messageBO.getSessionId(), new AssistantMessage("[ERROR] " + e.getMessage()));
            log.error("Connection error while generating response: {}", e.getMessage());
            throw new RuntimeException("Failed to connect to AI service", e);
        } catch (Exception e) {
            chatMemory.add(messageBO.getSessionId(), new AssistantMessage("[ERROR] " + e.getMessage()));
            log.error("Error generating response: {}", e.getMessage());
            throw new RuntimeException("AI service error", e);
        }
    }

    @Override
    public Flux<ServerSentEvent<String>> generateResponseStreamSSEFlux(ModelConfigDO modelConfigDO, MessageDTO messageBO) {
        // 1. 从外部获取或持久化中恢复ChatMemory
        ChatMemory chatMemory = sessionChatMemoryMap.computeIfAbsent(messageBO.getSessionId(), sessionId -> MessageWindowChatMemory.builder().chatMemoryRepository(chatMemoryRepository).maxMessages(messageBO.getMaxMessages() == null ? 10 : messageBO.getMaxMessages()).build());
        log.info("Generating response stream using OpenAI model: {}", modelConfigDO.getModelName());
        // 从缓存中获取模型客户端
        OpenAiChatModel model = this.getChatModelFromCache(modelConfigDO.getId().toString());
        ChatClient chatClient = ChatClient.builder(model)
                .defaultAdvisors(
                        PromptChatMemoryAdvisor.builder(chatMemory).conversationId(messageBO.getSessionId()).order(1).build(),
                        ChatModelStreamAdvisor.builder().chatModel(model).build(),
                        new CustomSessionUsageAdvisor(userSessionUsageService)
                )
                .build();
        Prompt prompt = this.buildChatModelRequest(chatMemory, messageBO);
//        log.info("generateResponseStream Generated prompt for OpenAI model: {}", prompt);
        try {
            ToolCallback[] toolCallbacks = agentPluginService.getToolCallbacks(messageBO.getSessionId());
            Map<String, Object> context = getContextMap(modelConfigDO, messageBO, toolCallbacks);
            // 构建ChatClient 注册MCP
            Flux<ChatResponse> chatResponseFlux = Optional.ofNullable(toolCallbacks)
                    .filter(callbacks -> callbacks.length > 0)
                    .map(callbacks -> chatClient.prompt(prompt)
                            .advisors(advisorSpec -> advisorSpec.params(context))
                            .toolCallbacks(callbacks)
                            .stream()
                            .chatResponse()
                    )
                    .orElseGet(() -> chatClient.prompt(prompt)
                            .advisors(advisorSpec -> advisorSpec.params(context))
                            .stream()
                            .chatResponse());
            return ChatFluxUtilSSe.returnChatFlux(chatResponseFlux, chatMemory, messageBO.getSessionId());
        } catch (ResourceAccessException e) {
            // TODO 此处细化异常 日后可能会针对网络异常进行重试等处理
            chatMemory.add(messageBO.getSessionId(), new AssistantMessage(e.getMessage()));
            log.error("Connection error while streaming response: {}", e.getMessage());
            throw new RuntimeException("Failed to connect to AI service", e);
        } catch (Exception e) {
            chatMemory.add(messageBO.getSessionId(), new AssistantMessage(e.getMessage()));
            log.error("Unexpected error while streaming response: {}", e.getMessage());
            throw new RuntimeException("AI service error", e);
        }
    }

    @NotNull
    private static Map<String, Object> getContextMap(ModelConfigDO modelConfigDO, MessageDTO messageBO, ToolCallback[] toolCallbacks) {
        Map<String, Object> context = new HashMap<>();
        context.put("userId", messageBO.getUserId());
        context.put("modelConfigDO", modelConfigDO);
        context.put("sessionId", messageBO.getSessionId());
        context.put("agentId", messageBO.getAgentId());
        if (toolCallbacks != null) {
            context.put("toolCallbacks", toolCallbacks);
        }
        if (StringUtils.hasText(messageBO.getUserApiKey())) {
            context.put("userApiKey", messageBO.getUserApiKey());
        }
        return context;
    }

    @Override
    public ChatClientType getClientType() {
        return ChatClientType.OpenAi;
    }

    private Prompt buildChatModelRequest(ChatMemory chatMemory, MessageDTO messageDTO) {
        List<Message> messages = new ArrayList<>(chatMemory.get(messageDTO.getSessionId()));
        messageDTO.getMessages().forEach(message -> {
            Map<String, Object> metadata = new HashMap<>();
            metadata.put("showFlag", message.getShowFlag());
            if (message.getRole().equalsIgnoreCase("user")) {
                // 转换消息为jsonarray
                String content1 = message.getContent();
                List<ContentDTO> contentDTOS = JSONArray.parseArray(content1, ContentDTO.class);
                for (ContentDTO content : contentDTOS) {
                    List<Media> mediaList = new ArrayList<>();
                    // 用户文本消息
                    String userPrompt = Optional.ofNullable(content).filter(c -> c.getType().equalsIgnoreCase("text")).map(ContentDTO::getText).orElse(DEFAULT_USER_MESSAGE);
                    // 图片&文件消息
                    if (Objects.requireNonNull(content).getType().equalsIgnoreCase("image_url")) {
                        mediaList.add(new Media(determineMimeTypeFromUrl(content.getUrl()), URI.create(content.getUrl())));
                    }
                    if (Objects.requireNonNull(content).getType().equalsIgnoreCase("file_url")) {
                        mediaList.add(new Media(determineMimeTypeFromUrl(content.getUrl()), URI.create(content.getUrl())));
                    }
                    // 文件
                    if (!ObjectUtils.isEmpty(message.getFile())) {
                        metadata.put("file", JSONArray.toJSONString(message.getFile()));
                    }
                    UserMessage userMessage = UserMessage.builder().text(Objects.requireNonNull(userPrompt)).media(mediaList).metadata(metadata).build();
                    messages.add(userMessage);

                }

            }
            if (message.getRole().equalsIgnoreCase("system")) {
                String content = ObjectUtils.isEmpty(message.getContent()) ? DEFAULT_SYSTEM_PROMPT : message.getContent();
                if (StringUtils.hasText(messageDTO.getUserPhone())) {
                    String userInfoSys = """
                            ## 固定参数信息 （注意：手机号只能从下面这个值中取，不能从上下文缓存中获取）
                            $_yumo_agent_user_phone_$=%s
                            """;
                    userInfoSys = String.format(userInfoSys, messageDTO.getUserPhone());
                    content = content + "\n" + userInfoSys;
                    log.info("userInfoSys:{},sessionId():{}", userInfoSys, messageDTO.getSessionId());
                }

                SystemMessage system = SystemMessage.builder()
                        .text(content).metadata(metadata).build();
                messages.add(system);

            }
        });

        return new Prompt(messages);
    }

    private OpenAiChatModel getChatModelFromCache(String modelId) {
        // 从缓存中获取模型客户端
        OpenAiChatModel client = aiClientModelCache.getOpenAiChatModel(modelId);
        if (client == null) {
            // 刷新缓存
            aiClientModelCache.init();
            client = aiClientModelCache.getOpenAiChatModel(modelId);
        }
        if (client == null) {
            log.error("Model not found in cache: {}", modelId);
            throw new RuntimeException("Model not found or not initialized");
        }
        return client;
    }

    /**
     * 根据 URL 后缀返回对应的 MIME 类型
     */
    private static MimeType determineMimeTypeFromUrl(String url) {
        if (url == null) {
            return MimeTypeUtils.ALL; // 默认值
        }
        // 1. 去掉 ? 和 # 之后的部分（查询参数和哈希）
        String lowerUrl = url.split("[?#]")[0].toLowerCase();

        if (lowerUrl.endsWith(".jpg") || lowerUrl.endsWith(".jpeg")) {
            return MimeTypeUtils.IMAGE_JPEG;
        } else if (lowerUrl.endsWith(".png")) {
            return MimeTypeUtils.IMAGE_PNG;
        } else if (lowerUrl.endsWith(".gif")) {
            return MimeTypeUtils.IMAGE_GIF;
        } else if (lowerUrl.endsWith(".txt")) {
            return MimeTypeUtils.TEXT_PLAIN;
        } else if (lowerUrl.endsWith(".json")) {
            return MimeTypeUtils.APPLICATION_JSON;
        } else if (lowerUrl.endsWith(".xml")) {
            return MimeTypeUtils.APPLICATION_XML;
        } else if (lowerUrl.endsWith(".html")) {
            return MimeTypeUtils.TEXT_HTML;
        }
        // 默认值
        return MimeTypeUtils.ALL;
    }
}
